from celery import shared_task
import os
# import fabric
# @shared_task(name='app.machine.add_node')
# def add_node(node_info):
#     connect_kwargs = {"password": node_info["password"]}
#     conn = fabric.Connection(host=node_info["hostip"],
#                              port=node_info["sshport"],
#                              user=node_info["username"],
#                              connect_kwargs=connect_kwargs,
#                              )
#     conn.config.sudo.password = node_info["password"]
#     try:
#         conn.open()
#         print("collection state {}".format(conn.is_connected))
#         conn.sudo('bash /home/{}/worker_node_add/node_add.sh'.format(node_info["username"]), hide=True, pty=False)
#     except Exception as e:
#         print(str(e))
#
#     conn.close()

def get_home_dir(user):
    if user == "root":
        return "/root"

    return "/home/{}".format(user)

@shared_task(name='app.machine.add_node')
def add_node(node_info):
    import paramiko
    conn_session = paramiko.SSHClient()
    conn_session.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    conn_session.connect(node_info["hostip"],
                         node_info["sshport"],
                         node_info["username"],
                         node_info["password"],
                         look_for_keys=False,
                         allow_agent=False)
    conn = conn_session.invoke_shell()  # to keep the session go on
    conn.keep_this = conn_session
    home_dir = get_home_dir(node_info["username"])
    log_remote_path = "{}/node_add.log".format(home_dir)
    stdin, stdout, stderr = conn_session.exec_command('sudo bash {}/worker_node_add/node_add.sh >{} 2>&1'.format(home_dir, log_remote_path), get_pty=True)
    stdin.write(node_info["password"] + '\n')
    stdin.flush()
    # for response_msg in stdout:
    #     print(response_msg.strip('\n'))
    # for response_msg in stderr:
    #     print(response_msg.strip('\n'))

    exit_code = stdout.channel.recv_exit_status()
    conn_session.close()
    print("exit code is {}".format(exit_code))

    #任务结束后将日志download到服务器上
    transport = paramiko.Transport((node_info["hostip"], node_info["sshport"]))
    transport.connect(username=node_info["username"], password=node_info["password"])
    sftp = paramiko.SFTPClient.from_transport(transport)

    dir_path = "/node_operation/log/{0}".format(node_info["hostip"])
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    local_path = "/node_operation/log/{0}/log-{0}".format(node_info["hostip"])
    sftp.get(log_remote_path, local_path)
    transport.close()

    return exit_code
